Description

Summary

Makes tensorflow, tensorflow-datasets, tensorflow-text, Jetstream and seqio optional dependencies in MaxText (Jetstream pulls tensorflow through Jetstream -> seqio_tensorflow-text -> tensorflow). Core pre-training lock files are now TensorFlow-free, while users requiring legacy TFDS pipelines or JetStream serving can install optional dependencies via --with-tf.

b/555209521

Changes

• Requirements & Lock Files:
  • Removed TensorFlow, TFDS, TF-Text, and SeqIO from base requirements and generated lock files (tpu-requirements.txt, cuda12-requirements.txt, decoupled-requirements.txt).
  • Added src/dependencies/extra_deps/tf_requirements.txt and jetstream_github_deps.txt for optional installation.
• Installation & Setup:
  • Added --with-tf flag (and WITH_TF env var) to install_pre_train_extra_deps.py to install optional TensorFlow and JetStream dependencies.
  • Added TF option to src/dependencies/scripts/setup.sh and docker_build_dependency_image.sh (defaults to false).
  • Added ARG TF=true to Dockerfiles because nightly image workflow does not pass "TF" in build arg yet, will keep TF in nightly images until we build separated workflow for TF and no-TF images.
• JetStream Stubs & Tests:
  • Updated gcloud_stub.py to return stubs when JetStream or its dependencies are absent, avoiding import crashes.
  • Added skip guards to JetStream-dependent tests (inference_microbenchmark_smoke_test.py, offline_engine_test.py). The following tests are skipped in the no-TF environment. Will follow up with a dedicated jetstream CI workflow for these tests similar to the vllm post-training CI workflow.

    * `tests/integration/smoke/inference_microbenchmark_smoke_test.py`
    * `tests/unit/offline_engine_test.py`
    * `tests/integration/decode_tests.py`
    * `tests/integration/generate_param_only_checkpoint_test.py`
    * `tests/integration/maxengine_test.py`
    * `tests/unit/maxengine_nnx_test.py`
    * `tests/unit/deepseek_decode_consistency_test.py`

• Added automatic JetStream checks to 42 end-to-end test scripts calling maxtext.inference.decode to install --with-tf if absent.
• Documentation: Updated install_maxtext.md and data_input_tfds.md with --with-tf instructions.

Tests

• github CI test pass (The failure in tests/unit/mmap_data_processing_test.py::GrainMmapNpyEvalConfigTest is at head and unrelated to this PR)
• Generated environments following this instruction for cases with or without TF. Confirmed it works as expected.

uv pip install -e .[tpu] --resolution=lowest
install_tpu_pre_train_extra_deps  # tested with and without --with-tf

• Tested generating tf and no-tf maxtext images, confirm reduced image size (local uncompressed size 5.95 GB -> 3.91 GB, 35% reduction)

bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=tpu MODE=stable TF=true LOCAL_IMAGE_NAME=maxtext_tf
bash src/dependencies/scripts/docker_build_dependency_image.sh DEVICE=tpu MODE=stable TF=false LOCAL_IMAGE_NAME=maxtext_no_tf

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.